"""
合同分块及入库测试
"""
import json
from pathlib import Path
from core.config import settings
from core.parse import ContractParser
from core.milvus_vector import MilvusVector
from loguru import logger


if __name__ == '__main__':

    # 读取合同数据
    contract_data = json.load(open(Path(__file__).parent.parent / r'data/资产转让合同.json', 'r', encoding='utf-8'))

    # 清洗及文本分块
    parser = ContractParser(
        chunk_size=settings.CHUNK_SIZE,
        chunk_overlap=settings.CHUNK_OVERLAP
    )

    logger.info("开始合同文本分块.")
    documents = parser.parse(contract_data)
    logger.info("结束合同文本分块.")

    for doc in documents:
        print(doc)

    # Embedding及入库milvus
    milvus = MilvusVector()
    contract_id = contract_data['fileId'].split('-')[-1]

    logger.info("开始文本块向量嵌入及入库.")
    milvus.create(collection_name=contract_id, documents=documents)
    logger.info("结束文本块向量嵌入及入库.")